from pymodbus.exceptions import ModbusException
from typing import List, Optional

from enums.decode import Decode
from enums.protocol_type import ProtocolType, TcpConfig, RtuConfig
from service.client.modbus_rtu_client import ModbusRtuClient
from service.client.modbus_rtu_over_tcp_client import ModbusRtuOverTcpClient
from service.log import log
from service.client.modbus_tcp_client import ModbusTcpClient


class ModbusClient:
    """
    Modbus TCP 客户端类，封装了常见的 Modbus 操作。
    """

    def __init__(self, protocol_type: ProtocolType, tcp_config: TcpConfig | None, rtu_config: RtuConfig | None):
        self.protocol_type = protocol_type
        if protocol_type == ProtocolType.TCP or protocol_type == ProtocolType.RTU_OVER_TCP:
            self.set_tcp_client(tcp_config)
        elif protocol_type == ProtocolType.RTU:
            self.set_rtu_client(rtu_config)

    def set_tcp_client(self, tcp_config: TcpConfig):
        if tcp_config.protocol_type == ProtocolType.TCP:
            self.client = ModbusTcpClient(tcp_config.host, tcp_config.port)
        elif tcp_config.protocol_type == ProtocolType.RTU_OVER_TCP:
            self.client = ModbusRtuOverTcpClient(tcp_config.host, tcp_config.port)

    def set_rtu_client(self, rtu_config: RtuConfig):
        self.client = ModbusRtuClient(rtu_config.port, rtu_config.baudrate, rtu_config.parity, rtu_config.stopbits,
                                      rtu_config.bytesize, rtu_config.timeout)

    def connect(self) -> bool:
        """
        连接到 Modbus 服务器。
        :return: 连接成功返回 True，否则返回 False
        """
        try:
            if self.client.connect():
                log.info(f"成功连接到 Modbus 服务器")
                return True
            else:
                log.error(f"无法连接到 Modbus 服务器")
                return False
        except Exception as e:
            log.error(f"连接 Modbus 服务器时发生异常: {e}")
            return False

    def disconnect(self):
        """
        断开与 Modbus 服务器的连接。
        """
        self.client.close()
        log.info(f"已断开与 Modbus 服务器的连接")

    def get_latest_message(self) -> list:
        return self.client.get_latest_message()

    def read_register(self, slave_id: int, register_type: int, start_address: int, end_address: int) -> List[int]:
        value_list = []
        if register_type == 1:
            value_list = self.read_coil_registers(slave_id, start_address, end_address - start_address + 1)
        elif register_type == 2:
            value_list = self.read_discrete_input_registers(slave_id, start_address, end_address - start_address + 1)
        elif register_type == 3:
            value_list = self.read_holding_registers(slave_id, start_address, end_address - start_address + 1)
        elif register_type == 4:
            value_list = self.read_input_registers(slave_id, start_address, end_address - start_address + 1)
        return value_list

    def read_coil_registers(self, slave_id: int, address: int, count: int) -> Optional[List[int]]:
        """
        读取线圈寄存器。
        :param slave_id: 从机地址
        :param address:  寄存器起始地址
        :param count:    要读取的寄存器数量
        :return:         读取到的寄存器值列表，如果失败则返回 None
        """
        try:
            # 按照120个寄存器分批读取
            MAX_REGISTERS_CNT = 120
            result = []
            for offset in range(0, count, MAX_REGISTERS_CNT):
                batch_size = min(MAX_REGISTERS_CNT, count - offset)
                current_address = address + offset

                try:
                    response = self.client.read_coils(
                        slave=slave_id,
                        address=current_address,
                        count=batch_size
                    )

                    if response.isError():
                        log.error(
                            f"读取线圈寄存器失败: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                        return result

                    result.extend([1 if bit else 0 for bit in response.bits])
                    log.debug(
                        f"成功读取线圈寄存器: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                except Exception as e:
                    log.error(f"读取线圈寄存器异常: {str(e)}")
                    return None
            return result
        except ModbusException as e:
            log.error(f"读取线圈寄存器时发生 Modbus 异常: {e}")
            return None

    def read_discrete_input_registers(self, slave_id: int, address: int, count: int) -> Optional[List[int]]:
        """
        读取离散寄存器。
        :param slave_id: 从机地址
        :param address:  寄存器起始地址
        :param count:    要读取的寄存器数量
        :return:         读取到的寄存器值列表，如果失败则返回 None
        """
        try:
            # 按照120个寄存器分批读取
            MAX_REGISTERS_CNT = 120
            result = []
            for offset in range(0, count, MAX_REGISTERS_CNT):
                batch_size = min(MAX_REGISTERS_CNT, count - offset)
                current_address = address + offset

                try:
                    response = self.client.read_discrete_inputs(
                        slave=slave_id,
                        address=current_address,
                        count=batch_size
                    )

                    if response.isError():
                        log.error(
                            f"读取离散寄存器失败: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                        return result

                    result.extend([1 if bit else 0 for bit in response.bits])
                    log.debug(
                        f"成功读取离散寄存器: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                except Exception as e:
                    log.error(f"读取离散寄存器异常: {str(e)}")
                    return None
            return result
        except ModbusException as e:
            log.error(f"读取离散寄存器时发生 Modbus 异常: {e}")
            return None

    def read_holding_registers(self, slave_id: int, address: int, count: int) -> Optional[List[int]]:
        """
        读取保持寄存器。
        :param slave_id: 从机地址
        :param address:  寄存器起始地址
        :param count:    要读取的寄存器数量
        :return:         读取到的寄存器值列表，如果失败则返回 None
        """
        try:
            # 按照120个寄存器分批读取
            MAX_REGISTERS_CNT = 120
            result = []
            for offset in range(0, count, MAX_REGISTERS_CNT):
                batch_size = min(MAX_REGISTERS_CNT, count - offset)
                current_address = address + offset

                try:
                    response = self.client.read_holding_registers(
                        slave=slave_id,
                        address=current_address,
                        count=batch_size
                    )

                    if response.isError():
                        log.error(f"读取失败: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                        return result

                    result.extend(response.registers)
                    log.debug(f"成功读取: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                except Exception as e:
                    log.error(f"读取异常: {str(e)}")
                    return None
            return result
        except ModbusException as e:
            log.error(f"读取保持寄存器时发生 Modbus 异常: {e}")
            return None
        except Exception as e:
            log.error(f"读取保持寄存器时发生未知异常: {e}")
            return None

    def read_input_registers(self, slave_id: int, address: int, count: int) -> Optional[List[int]]:
        """
        读取输入寄存器。
        :param slave_id: 从机地址
        :param address:  寄存器起始地址
        :param count:    要读取的寄存器数量
        :return:         读取到的寄存器值列表，如果失败则返回 None
        """
        try:
            # 按照120个寄存器分批读取
            MAX_REGISTERS_CNT = 120
            result = []
            for offset in range(0, count, MAX_REGISTERS_CNT):
                batch_size = min(MAX_REGISTERS_CNT, count - offset)
                current_address = address + offset

                try:
                    response = self.client.read_input_registers(
                        slave=slave_id,
                        address=current_address,
                        count=batch_size
                    )

                    if response.isError():
                        log.error(f"读取失败: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                        return result

                    result.extend(response.registers)
                    log.debug(f"成功读取: 从机{slave_id} 地址{current_address}-{current_address + batch_size - 1}")
                except Exception as e:
                    log.error(f"读取输入寄存器时出错: {str(e)}")
                    return None
            return result
        except ModbusException as e:
            log.error(f"读取输入寄存器时发生 Modbus 异常: {e}")
            return None

    def write_single_coil_register(self, slave_id: int, start_address: int, value_list: List[bool]) -> bool:
        """
        写入单个线圈寄存器。
        参数:
            slave_id: 从机地址 (1-247)
            start_address: 寄存器地址 (0-65535)
            value_list: 要写入的值列表
        返回:
            bool: 成功返回True，失败返回False
        """
        try:
            for index, value in enumerate(value_list):
                response = self.client.write_coil(
                    address=start_address + index,
                    value=value,
                    slave=slave_id
                )
                if response.isError():
                    log.error(f"写入线圈寄存器时出错: {response}")
                    return False
            return True
        except ModbusException as e:
            log.error(f"写入线圈寄存器时发生 Modbus 异常: {e}")
            return False
        except Exception as e:
            log.error(f"写入线圈寄存器时发生未知异常: {e}")
            return False

    def write_coil_registers(self, slave_id: int, start_address: int, value_list: List[bool]) -> bool:
        """
        批量写入线圈寄存器（支持多寄存器）
        参数:
            slave_id: 从机地址 (1-247)
            start_address: 起始寄存器地址 (0-65535)
            values: 要写入的值列表
        返回:
            bool: 成功返回True，失败返回False
        """
        try:
            response = self.client.write_coils(
                address=start_address,
                values=value_list,
                slave=slave_id
            )
            return not response.isError()
        except ModbusException as e:
            log.error(f"写入线圈寄存器时发生 Modbus 异常: {e}")
            return False
        except Exception as e:
            log.error(f"写入线圈寄存器时发生未知异常: {e}")
            return False

    def write_single_holding_register(
            self,
            slave_id: int,
            start_address: int,
            data_type: str,
            value_list: List[int],
    ) -> bool:
        """
        单个写入保持寄存器（支持多种数据类型和循环写入）
        参数:
            slave_id: 从机地址 (1-247)
            start_address: 寄存器地址 (0-65535)
            value_list: 要写入的值（支持整数/浮点数）
            data_type: 数据类型（如 "float_abcd"）
        返回:
            bool: 成功返回True，失败返回False
        """
        global response
        try:
            byteorder = Decode.byteorder_mapping.get(data_type, ">H")
            current_address = start_address

            for value in value_list:
                packed = Decode.pack_value(byteorder, value)

                # 处理单寄存器写入（16位）
                if len(packed) <= 2:
                    register_value = int.from_bytes(packed, byteorder="big")
                    response = self.client.write_register(
                        address=current_address,
                        value=register_value,
                        slave=slave_id
                    )
                    current_address += 1
                else:
                    # 处理多寄存器数据类型（如32位浮点数）
                    registers = []
                    for i in range(0, len(packed), 2):
                        chunk = packed[i:i + 2]
                        registers.append(int.from_bytes(chunk, byteorder="big"))

                    # 循环写入拆分后的寄存器
                    for reg_value in registers:
                        response = self.client.write_register(
                            address=current_address,
                            value=reg_value,
                            slave=slave_id
                        )
                        current_address += 1
                        if response.isError():
                            return False

                if response.isError():
                    log.error(f"写入保持寄存器时出错: {response}")
                    return False

            return True
        except ModbusException as e:
            log.error(f"(06功能码)写入保持寄存器异常: {e}")
            return False
        except Exception as e:
            log.error(f"(06功能码)未知异常: {e}")
            return False

    def write_holding_registers(
            self,
            slave_id: int,
            start_address: int,
            data_type: str,
            value_list: List[int],
    ) -> bool:
        """
        批量写入保持寄存器（支持多寄存器和多种数据类型）
        参数:
            slave_id: 从机地址 (1-247)
            start_address: 起始寄存器地址 (0-65535)
            values: 要写入的值列表（支持整数/浮点数）
            data_type: 数据类型（如 "float_abcd"）
        返回:
            bool: 成功返回True，失败返回False
        """
        try:
            # 1. 根据数据类型转换值为寄存器格式
            byteorder = Decode.byteorder_mapping.get(data_type, ">H")
            registers = []
            for value in value_list:
                packed = Decode.pack_value(byteorder, value)
                # 拆分bytes为16位寄存器值（兼容多寄存器）
                for i in range(0, len(packed), 2):
                    chunk = packed[i:i + 2]
                    registers.append(int.from_bytes(chunk, byteorder="big"))

            # 2. 批量写入寄存器
            response = self.client.write_registers(
                address=start_address,
                values=registers,
                slave=slave_id
            )
            return not response.isError()
        except ModbusException as e:
            log.error(f"(16功能码)写入保持寄存器时发生 Modbus 异常: {e}")
            return False
        except Exception as e:
            log.error(f"(16功能码)写入保持寄存器时发生未知异常: {e}")
            return False

    def __enter__(self):
        """支持上下文管理器，用于 with 语句"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持上下文管理器，用于 with 语句"""
        self.disconnect()


if __name__ == "__main__":
    # 使用示例
    client = ModbusClient(host="10.8.0.126", port=10502)
    registers = client.read_holding_registers(slave_id=1, address=0, count=10)
    if registers:
        log.info(f"读取到的寄存器值: {registers}")
    else:
        log.error("读取寄存器失败")
